#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
逆运动学数据记录器
用于记录和重放逆运动学计算所需的数据
"""

import os
import time
import numpy as np
import pickle
import json
from datetime import datetime

class IKDataRecorder:
    def __init__(self, data_dir=None, auto_save_interval=100):
        """
        初始化数据记录器
        
        参数:
            data_dir: 数据保存目录，如果为None则使用默认目录
            auto_save_interval: 自动保存的间隔（记录数据条数），设为0禁用自动保存
        """
        # 设置数据保存目录
        if data_dir is None:
            # 默认保存在项目根目录下的data/ik_data目录
            current_dir = os.path.dirname(os.path.abspath(__file__))
            project_root = os.path.dirname(os.path.dirname(current_dir))
            self.data_dir = os.path.join(project_root, 'data', 'ik_data')
        else:
            self.data_dir = data_dir
            
        # 确保目录存在
        os.makedirs(self.data_dir, exist_ok=True)
        
        # 初始化数据列表
        self.data_list = []
        
        # 当前记录会话的文件名
        self.current_session_file = None
        
        # 自动保存设置
        self.auto_save_interval = auto_save_interval
        self.last_save_count = 0
        
        print(f"IK数据记录器初始化完成，数据将保存到: {self.data_dir}")
        if auto_save_interval > 0:
            print(f"已启用自动保存功能，每记录{auto_save_interval}条数据自动保存一次")
    
    def start_new_session(self):
        """
        开始新的记录会话
        """
        # 生成会话文件名，使用时间戳
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.current_session_file = os.path.join(self.data_dir, f"ik_data_{timestamp}.pkl")
        
        # 清空数据列表
        self.data_list = []
        
        print(f"开始新的IK数据记录会话: {self.current_session_file}")
        return self.current_session_file
    
    def record(self, left_wrist, right_wrist, current_lr_arm_q=None, current_lr_arm_dq=None, sol_q=None):
        """
        记录逆运动学数据
        
        参数:
            left_wrist: 左手腕位姿
            right_wrist: 右手腕位姿
            current_lr_arm_q: 当前左右臂关节角度
            current_lr_arm_dq: 当前左右臂关节速度
            sol_q: 逆运动学解算得到的关节角度
        """
        # 确保当前会话文件名已设置
        if self.current_session_file is None:
            self.start_new_session()
            
        # 创建数据记录
        data_record = {
            'timestamp': time.time(),
            'left_wrist': left_wrist,
            'right_wrist': right_wrist,
            'current_lr_arm_q': current_lr_arm_q,
            'current_lr_arm_dq': current_lr_arm_dq,
            'sol_q': sol_q  # 添加逆运动学解算结果
        }
        
        # 添加到数据列表
        self.data_list.append(data_record)
        
        # 检查是否需要自动保存
        if self.auto_save_interval > 0 and len(self.data_list) - self.last_save_count >= self.auto_save_interval:
            self.save()
            self.last_save_count = len(self.data_list)
            print(f"[自动保存] 已保存{len(self.data_list)}条IK数据到{self.current_session_file}")
    
    def save(self):
        """
        保存当前会话数据到文件（可用于自动保存）
        
        返回:
            保存的文件路径
        """
        return self.save_session()
        
    def save_session(self):
        """
        保存当前会话数据到文件
        
        返回:
            保存的文件路径
        """
        # 确保有数据要保存
        if not self.data_list:
            print("没有数据需要保存")
            return None
        
        # 确保会话已经开始
        if self.current_session_file is None:
            self.start_new_session()
        
        # 保存数据到pickle文件
        with open(self.current_session_file, 'wb') as f:
            pickle.dump(self.data_list, f)
        
        # 保存数据摘要到JSON文件（便于查看）
        summary_file = self.current_session_file.replace('.pkl', '_summary.json')
        summary = {
            'total_records': len(self.data_list),
            'start_time': datetime.fromtimestamp(self.data_list[0]['timestamp']).strftime("%Y-%m-%d %H:%M:%S") if self.data_list else None,
            'end_time': datetime.fromtimestamp(self.data_list[-1]['timestamp']).strftime("%Y-%m-%d %H:%M:%S") if self.data_list else None,
            'duration_seconds': self.data_list[-1]['timestamp'] - self.data_list[0]['timestamp'] if len(self.data_list) > 1 else 0
        }
        
        with open(summary_file, 'w', encoding='utf-8') as f:
            json.dump(summary, f, ensure_ascii=False, indent=4)
        
        print(f"保存了 {len(self.data_list)} 条IK数据记录到: {self.current_session_file}")
        print(f"数据摘要保存到: {summary_file}")
        
        return self.current_session_file
    
    def load_session(self, file_path=None):
        """
        加载会话数据
        
        参数:
            file_path: 要加载的文件路径，如果为None则使用当前会话文件
            
        返回:
            加载的数据列表
        """
        if file_path is None:
            file_path = self.current_session_file
        
        if file_path is None or not os.path.exists(file_path):
            print(f"文件不存在: {file_path}")
            return []
        
        try:
            with open(file_path, 'rb') as f:
                self.data_list = pickle.load(f)
            
            print(f"从 {file_path} 加载了 {len(self.data_list)} 条IK数据记录")
            return self.data_list
        except Exception as e:
            print(f"加载数据失败: {e}")
            return []
    
    def get_latest_session_file(self):
        """
        获取最新的会话文件路径
        
        返回:
            最新的会话文件路径，如果没有则返回None
        """
        # 获取所有pkl文件
        pkl_files = [os.path.join(self.data_dir, f) for f in os.listdir(self.data_dir) 
                    if f.endswith('.pkl') and f.startswith('ik_data_')]
        
        if not pkl_files:
            return None
        
        # 按修改时间排序，返回最新的
        return max(pkl_files, key=os.path.getmtime)
